// Copyright (C) 2023 Computer Vision Lab, Electrical Engineering, 
// Indian Institute of Science, Bengaluru, India.
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//     * Neither the name of Indian Institute of Science nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
// DEALINGS IN THE SOFTWARE.
// 
// 
// This file is a part of the implementation for the paper:
// Chitturi Sidhatha, Lalit Manam and Venu Madhav Govindu, 
// Adaptive Annealing for Robust Geometric Estimation,
// Proceedings of the Computer Vision and Pattern Recognition 2023
// 
// This code base is adapted from the code release of Fast Global Registration,
// Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun, ECCV 2016

#include <stdio.h>
#include <chrono>
#include "app.h"
#include "mex.h"

#define IS_REAL_2D_FULL_DOUBLE(P) (!mxIsComplex(P) && mxGetNumberOfDimensions(P)==2 && !mxIsSparse(P) && mxIsDouble(P))
#define IS_REAL_2D_OR_3D_FULL_DOUBLE(P) (!mxIsComplex(P) && (mxGetNumberOfDimensions(P)==2 || mxGetNumberOfDimensions(P)==3) && !mxIsEmpty(P) && !mxIsSparse(P) && mxIsDouble(P))
#define IS_REAL_SCALAR(P) (IS_REAL_2D_FULL_DOUBLE(P) && mxGetNumberOfElements(P)==1)
#define IS_LOGICAL_SCALAR(P) (mxIsLogicalScalar(P))


#define P1_MATCHED_IN prhs[0]
#define P2_MATCHED_IN prhs[1]
#define THRES_IN prhs[2]
#define SIG_MIN_IN prhs[3]
#define EIG_MIN_THRES_IN prhs[4]
#define ISPERTURB_IN prhs[5]
#define FIXED_ANNEAL_IN prhs[6]
#define ANNEAL_FACTOR_IN prhs[7]
#define ISUMEYA_IN prhs[8]
#define SIG_MAX_IN prhs[9]
#define STOCH_SAMPLING_IN prhs[10]

#define MOTION_OUT plhs[0]
#define N_RUNS plhs[1]
#define N_TIMES_INNER plhs[2]
#define TIME_TAKEN plhs[3]
#define SIG_STOP plhs[4]
#define ITER plhs[5]
#define N_PERTURBS plhs[6]
// #define PERTURB_ROTANGLE plhs[7]
// #define PERTURB_TRANS plhs[8]

using namespace gom;

void mexFunction( int nlhs, mxArray *plhs[],
        		  int nrhs, const mxArray*prhs[] )
{
    double thres,sig_min,eig_min_thres,anneal_factor,sig_max;
    bool isperturb, fixed_anneal,isumeya,stoch_sampling;

    /* Some checks */
    // if(nrhs!=2)
    //     mexErrMsgTxt("FastGlobalRegistration(points_01, points_02)\n");
    // if(mxGetN(prhs[0])!=3 || mxGetN(prhs[1])!=3)
    //     mexErrMsgTxt("Points must have three columns\n");

    /* Some checks */
    if(nrhs < 2) mexErrMsgTxt("At least two input arguments are required.");
    else if(nrhs > 11) mexErrMsgTxt("Too many input arguments.");
    else if(nlhs < 1) mexErrMsgTxt("At least one output argument is required.");
    else if(nlhs > 7) mexErrMsgTxt("Too many output arguments.");

    if(!IS_REAL_2D_FULL_DOUBLE(P1_MATCHED_IN) || mxIsEmpty(P1_MATCHED_IN) || mxGetN(P1_MATCHED_IN)!=3)
        mexErrMsgTxt("Point set 1 must be real full array of size Nx3.");
    if(!IS_REAL_2D_FULL_DOUBLE(P2_MATCHED_IN) || mxIsEmpty(P2_MATCHED_IN) || mxGetN(P2_MATCHED_IN)!=3)
        mexErrMsgTxt("Point set 2 must be real full array of size 3xN.");
    if(mxGetM(P1_MATCHED_IN) != mxGetM(P2_MATCHED_IN))
        mexErrMsgTxt("Number of matched points must be same in both sets.");
    //THRES_IN
    if(nrhs < 3 || mxIsEmpty(THRES_IN))
        thres = 0.01;
    else {
        if(!IS_REAL_SCALAR(THRES_IN) || mxGetScalar(THRES_IN) < 0.0)
            mexErrMsgTxt("Threshold thres should be a positive real number.");
        else
            thres = mxGetScalar(THRES_IN);
    }
    //SIG_MIN_IN
    if(nrhs < 4 || mxIsEmpty(SIG_MIN_IN))
        sig_min = 0.1;
    else {
        if(!IS_REAL_SCALAR(SIG_MIN_IN) || mxGetScalar(SIG_MIN_IN) < 0.0)
            mexErrMsgTxt("Minimum sigma should be a positive real number.");
        else
            sig_min = mxGetScalar(SIG_MIN_IN);
    }    
    //EIG_MIN_THRES_IN
    if(nrhs < 5 || mxIsEmpty(EIG_MIN_THRES_IN))
        eig_min_thres = 0.0d;
    else {
        eig_min_thres = mxGetScalar(EIG_MIN_THRES_IN);
    }        
    //Decide whether perturb is present or not
    if(nrhs < 6 || mxIsEmpty(ISPERTURB_IN))
        isperturb = true;
    else {
        if(!IS_LOGICAL_SCALAR(ISPERTURB_IN))
            mexErrMsgTxt("isperturb must be a logical scalar.");
        else
            isperturb = mxIsLogicalScalarTrue(ISPERTURB_IN);
    }    
    //Decide if fixed_anneal flag is on/off
    if(nrhs < 7 || mxIsEmpty(FIXED_ANNEAL_IN))
        fixed_anneal = false;
    else {
        if(!IS_LOGICAL_SCALAR(FIXED_ANNEAL_IN))
            mexErrMsgTxt("fixed_anneal must be a logical scalar.");
        else
            fixed_anneal = mxIsLogicalScalarTrue(FIXED_ANNEAL_IN);
    }       
    //Annealing factor if fixed anneal is true
    if(nrhs < 8 || mxIsEmpty(ANNEAL_FACTOR_IN))
        anneal_factor = 1.4d;
    else {
        anneal_factor = mxGetScalar(ANNEAL_FACTOR_IN);
    }        
    //Decide if umeyama method should be used or SE3Reg should be used. Currently  only Umeyama method is functional.
    if(nrhs < 9 || mxIsEmpty(ISUMEYA_IN))
        isumeya = true;
    else {
        if(!IS_LOGICAL_SCALAR(ISUMEYA_IN))
            mexErrMsgTxt("isumeya must be a logical scalar.");
        else
            isumeya = mxIsLogicalScalarTrue(ISUMEYA_IN);
    }           
    //Maximum Sigma to be used in the first iteration
    if(nrhs < 10 || mxIsEmpty(SIG_MAX_IN))
        sig_max = 1000.0d;
    else {
        sig_max = mxGetScalar(SIG_MAX_IN);
    }              
    //Decide if stochastic sampling should be done or not.
    if(nrhs < 11 || mxIsEmpty(STOCH_SAMPLING_IN))
        stoch_sampling = false;
    else {
        if(!IS_LOGICAL_SCALAR(STOCH_SAMPLING_IN))
            mexErrMsgTxt("stoch_sampling must be a logical scalar.");
        else
            stoch_sampling = mxIsLogicalScalarTrue(STOCH_SAMPLING_IN);
    }          



    /* Input points */
    Eigen::Map<Eigen::Array<double,Eigen::Dynamic,3> >
            points_01((double*)mxGetData(prhs[0]),mxGetM(prhs[0]),mxGetN(prhs[0]));
    Eigen::Map<Eigen::Array<double,Eigen::Dynamic,3> >
            points_02((double*)mxGetData(prhs[1]),mxGetM(prhs[1]),mxGetN(prhs[1]));

    Eigen::MatrixXd points_11, points_12;
    points_11 = Eigen::Map<Eigen::MatrixXd >((double*)mxGetData(prhs[0]),mxGetM(prhs[0]),mxGetN(prhs[0])); // Rows are the number of points, cols are the dimension of the points.
    points_12 = Eigen::Map<Eigen::MatrixXd >((double*)mxGetData(prhs[1]),mxGetM(prhs[1]),mxGetN(prhs[1]));
   
    /* More checks */
    if(points_01.rows()!=points_02.rows())
        mexErrMsgTxt("Number of matched correspondences are not the same\n");

    /* Call the actual code */
	CApp app(stoch_sampling,sig_max,isumeya,anneal_factor,fixed_anneal,isperturb,eig_min_thres,sig_min,thres);
    app.LoadPointCloud(points_11, points_12);
	app.SetMatches();
    app.RobustUmeyamaPairwise();
	Eigen::Matrix4d transf = app.GetOutputTrans();

    /* Return the transformation */
    plhs[0] = mxCreateDoubleMatrix(4,4, mxREAL);
    double * pt = mxGetPr(plhs[0]);
    for(std::size_t ii=0; ii<4; ++ii)
        for(std::size_t jj=0; jj<4; ++jj)
            *(pt++) = transf(jj,ii);

    /* Return number of times umeyama method is run, n_runs*/
    plhs[1] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double * nruns = mxGetPr(plhs[1]); 
    *nruns = (double)app.n_runs_;

    /* Return number of times annealing by a fixed factor is done, n_times_inner*/
    plhs[2] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double *ntimesinner = mxGetPr(plhs[2]);
    *ntimesinner = (double)app.n_times_inner_;

    /* Return timing information */
    plhs[3] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double *time_net = mxGetPr(plhs[3]);
//    time_net[0] = ((double)rs.timeOptimize_)/CLOCKS_PER_SEC;
    *time_net = app.t_gom_;            

    /* Return stopping sigma */
    plhs[4] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double *sigstop = mxGetPr(plhs[4]);
    *sigstop = app.sig_stop_;

    /* Return number of Iterations */
    plhs[5] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double *iter = mxGetPr(plhs[5]);
    *iter = (double)app.n_iter_;

    /* Return number of perturbs */
    plhs[6] = mxCreateDoubleMatrix(1, 1, mxREAL);
    double *nperturbs = mxGetPr(plhs[6]);
    *nperturbs = (double)app.n_perturbs_;

}
